package com.xialuo.study.nodes;

import com.xialuo.study.state.ReactState;
import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.action.AsyncNodeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * 处理LLM交互的节点
 */
@Slf4j
public class LlmNode implements AsyncNodeAction<ReactState> {
    private final ChatClient chatClient;
    private final String promptTemplate;

    public LlmNode(ChatClient chatClient, String promptTemplate) {
        this.chatClient = chatClient;
        this.promptTemplate = promptTemplate;
    }

    @Override
    public CompletableFuture<Map<String, Object>> apply(ReactState state) {
        List<Message> messages = (List<Message>) state
                .value("messages")
                .get();

        // 如果第一条消息来自用户且我们有提示模板，
        // 在消息前添加一个系统消息
        if (messages.size() == 1 && messages.get(0) instanceof UserMessage && promptTemplate != null && !promptTemplate.isEmpty()) {
            List<Message> newMessages = new ArrayList<>();
            newMessages.add(new SystemMessage(promptTemplate));
            newMessages.addAll(messages);
            messages = newMessages;
        }

        // 从聊天客户端获取响应
        ChatResponse response = chatClient
                .prompt()
                .messages(messages)
                .call()
                .chatResponse();

        // 将助手的响应添加到消息列表中
        messages.add(response
                .getResult()
                .getOutput());
        log.debug("LLM Response: {}", response.getResult().getOutput());


        // 返回更新后的状态
        Map<String, Object> newState = new HashMap<>();
        newState.put("messages", messages);
        return CompletableFuture.completedFuture(newState);
    }
}
